#include <vector>
#include <iomanip> 
#include <iostream>
#include <chrono>
#include "matrix_operations_gpu.cu"
#include "matrix_operations_cpu.h"
#include "matrix_operations_wmma.cu"
#include "simple_fnn.h"
#include "relu.h"
#include "sigmoid.h"
#include "loss_function.cu"

bool areMatricesEqual(const std::vector<std::vector<float>>& A, const std::vector<std::vector<float>>& B, float tolerance = 1e-6) {
    if (A.size() != B.size() || A[0].size() != B[0].size()) {
        return false;
    }

    for (size_t i = 0; i < A.size(); ++i) {
        for (size_t j = 0; j < A[i].size(); ++j) {
            if (std::abs(A[i][j] - B[i][j]) > tolerance) {
                return false;
            }
        }
    }
    return true;
}

void printMatrix(const std::vector<std::vector<float>>& matrix) {
    for (const auto& row : matrix) {
        for (float elem : row) {
            std::cout << std::setw(10) << elem << " ";
        }
        std::cout << std::endl;
    }
}

int main(){
    // Example matrices
    auto startprog = std::chrono::high_resolution_clock::now();
    const int size = 8192;  // Set the size of the matrix to 16x16
    std::vector<std::vector<float>> A(size, std::vector<float>(size));
    std::vector<std::vector<float>> B(size, std::vector<float>(size));

    // Fill matrices with some values
    for (int i = 0; i < size; ++i) {
        for (int j = 0; j < size; ++j) {
            A[i][j] = 0.1;  // Simple pattern for easy verification
            B[i][j] = 0.1;  
            
        }
    }

    MatrixOperationsCPU cpu;
    MatrixOperationsGPU gpu;
    MatrixOperationsWMMA wmma;

    // Performance testing
    // auto startCPU = std::chrono::high_resolution_clock::now();
    // auto resultCPU = cpu.multiply(A, B);
    // auto endCPU = std::chrono::high_resolution_clock::now();
    // auto resultGPU1 = gpu.multiply(A, B);
    // auto startGPU = std::chrono::high_resolution_clock::now();
    // auto resultGPU = gpu.multiply(A, B);
    // auto endGPU = std::chrono::high_resolution_clock::now();
    // // auto resultWMM1 = wmma.multiply(A, B);
    auto startWMMA = std::chrono::high_resolution_clock::now();
    auto resultWMMA = wmma.multiply(A, B);
    auto endWMMA = std::chrono::high_resolution_clock::now();

    // std::chrono::duration<double, std::milli> cpuTime = endCPU - startCPU;
    // std::chrono::duration<double, std::milli> gpuTime = endGPU - startGPU;
    std::chrono::duration<double, std::milli> wmmaTime = endWMMA - startWMMA;

    // std::cout << "CPU Result:" << std::endl;
    // printMatrix(resultCPU);
    // std::cout << "GPU Result:" << std::endl;
    // printMatrix(resultGPU);

    // std::cout << "CPU Time: " << cpuTime.count() << " ms" << std::endl;
    // std::cout << "GPU Time: " << gpuTime.count() << " ms" << std::endl;
    std::cout << "WMMA Time: " << wmmaTime.count() << " ms" << std::endl;
    // std::cout << "Matrix A:" << std::endl;
    // printMatrix(A);

    // std::cout << "Matrix B:" << std::endl;
    // printMatrix(B);

    // // Compute on CPU and GPU
    // auto resultCPU = cpu.multiply(A, B);
    // auto resultGPU = wmma.multiply(A, B);  // This should be your GPU multiplication function

    // std::cout << "GPU Result:" << std::endl;
    // printMatrix(resultGPU);
    // std::cout << "CPU Result:" << std::endl;
    // printMatrix(resultCPU);

    // std::cout << "WMMA Result:" << std::endl;
    // printMatrix(resultWMMA);

    // Check if results are the same
    // if (areMatricesEqual(resultGPU, resultWMMA)) {
    //     std::cout << "Test passed!" << std::endl;
    // } else {
    //     std::cout << "Test failed!" << std::endl;
    // }
    auto endprog = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double, std::milli> progTime = endprog - startprog;
    std::cout << "prog Time: " << progTime.count() << " ms" << std::endl;
}